package com.stone.inventory.sharding.infra.hash;

import org.apache.commons.lang3.StringUtils;

import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 基于Ketama算法的一致性hash
 *
 * @author Mr_wenpan@163.com 2021/8/17 11:03 下午
 */
public class KetamaHash {

    /**
     * 一致性节点分布（存放虚拟节点，key: 虚拟节点的hash值，value：虚拟节点）
     */
    private final TreeMap<Long, String> nodeMap = new TreeMap<>();

    /**
     * 虚拟节点与真实节点映射(key: 虚拟节点，value: 真实节点)
     */
    private final Map<String, String> virtualMappingActual = new ConcurrentHashMap<>();

    private final ReentrantLock lock = new ReentrantLock();

    public KetamaHash(List<String> actualNodes) {
        init(actualNodes, DEFAULT_VIRTUAL_NODE_COUNT);
    }

    public KetamaHash(List<String> actualNodes, Integer virtualNodeCount) {
        init(actualNodes, virtualNodeCount == null ? DEFAULT_VIRTUAL_NODE_COUNT : virtualNodeCount);
    }

    /**
     * 初始化
     *
     * @param actualNodes      实际节点列表
     * @param virtualNodeCount 虚拟节点个数
     */
    private void init(List<String> actualNodes, Integer virtualNodeCount) {

        // 为每个真实节点虚拟出x个虚拟节点
        for (String actualNode : actualNodes) {
            for (int i = 0; i < virtualNodeCount; i++) {
                String virtualNode = actualNode + SYMBOL + i;
                // 使用Ketama算法求hash值
                long hash = Ketama.hash(virtualNode);
                // 保存虚拟节点
                nodeMap.put(hash, virtualNode);
                // 保存虚拟节点到真实节点的映射
                virtualMappingActual.put(virtualNode, actualNode);
            }
        }
    }

    /**
     * 获取真实节点
     *
     * @param key key
     * @return 真实节点
     */
    public String getActualNode(String key) {
        // 先求key的hash值
        long hash = Ketama.hash(key);
        // 通过key的hash值获取虚拟节点hash值
        Long index = nodeMap.tailMap(hash).firstKey();
        // 获取虚拟节点
        String virtualNode = nodeMap.get(index);
        // 通过虚拟节点获取真实节点
        return virtualMappingActual.get(virtualNode);
    }

    /**
     * 动态新增物理节点
     *
     * @param actualNodes      真实节点
     * @param virtualNodeCount 每个真实节点对应的虚拟节点个数
     * @author Mr_wenpan@163.com 2021/8/28 3:55 下午
     */
    public void addNodes(List<String> actualNodes, Integer virtualNodeCount) {
        // 新增节点基本没有并发，可以考虑同步，且ConcurrentHashMap允许一边遍历一边修改
        if (lock.tryLock()) {
            try {
                init(actualNodes, virtualNodeCount);
            } finally {
                lock.unlock();
            }
        }
    }

    /**
     * 动态减少真实物理节点
     *
     * @param actualNodes 要移除的真实节点
     * @author Mr_wenpan@163.com 2021/8/28 3:56 下午
     */
    public void removeNodes(List<String> actualNodes) {
        if (actualNodes == null || actualNodes.size() < 1) {
            return;
        }
        HashSet<String> virtualKeys = new HashSet<>();

        if (lock.tryLock()) {
            try {

                // 根据真实节点找到所有的要删除的虚拟节点(即，通过value找key)
                for (String actualNode : actualNodes) {
                    if (StringUtils.isBlank(actualNode)) {
                        continue;
                    }
                    for (Map.Entry<String, String> entry : virtualMappingActual.entrySet()) {
                        if (actualNode.equals(entry.getValue())) {
                            virtualKeys.add(entry.getKey());
                        }

                    }
                }

                // 通过虚拟节点在nodeMap中寻找对应的key(即，通过value找key)
                HashSet<Long> nodeKeys = new HashSet<>();
                for (String key : virtualKeys) {
                    for (Map.Entry<Long, String> entry : nodeMap.entrySet()) {
                        if (entry.getValue().equals(key)) {
                            nodeKeys.add(entry.getKey());
                        }
                    }
                }

                // 删除虚拟节点
                removeAll(nodeMap, nodeKeys);
                // 删除虚拟节点和真实节点的映射
                removeAll(virtualMappingActual, virtualKeys);
            } finally {
                lock.unlock();
            }
        }

    }

    private <K, V> void removeAll(Map<K, V> map, HashSet<K> keys) {
        if (keys == null || keys.size() < 1) {
            return;
        }
        for (K key : keys) {
            map.remove(key);
        }
    }

    /**
     * 默认的虚拟节点个数
     */
    private final Integer DEFAULT_VIRTUAL_NODE_COUNT = 160;

    /**
     * 符号
     */
    private final String SYMBOL = "_VN_";
}